"""
@Filename       : kg_interest_measure.py
@Create Time    : 2022/4/5 15:30
@Author         : Rylynn
@Description    : 

"""
import torch
import torch as th
import torch.nn as nn
import dgl

from model.kg.kg_encoder import MeanKGPooling


class DicreteKgInterestMeasure(nn.Module):
    def __init__(self, content_dict, user_dict):
        super(DicreteKgInterestMeasure, self).__init__()
        self.content_dict = content_dict
        self.user_dict = user_dict
        self.score_dict = None

    def forward(self, batch_content):
        score = torch.FloatTensor(len(batch_content), len(self.user_dict))
        batch_content_entities = [set(self.content_dict[c]) for c in batch_content]
        for content_entities in batch_content_entities:
            ...

    def calculate_static_score(self):
        ...




class ContinuousKgInterestMeasure(nn.Module):
    def __init__(self, content_dict, user_dict, entity_embedding, kg_encoder):
        super(ContinuousKgInterestMeasure, self).__init__()
        self.entity_embedding = entity_embedding
        self.content_dict = content_dict
        self.user_dict = user_dict
        self.kg_encoder = kg_encoder
        self.attention = nn.Parameter()

    def forward(self, batch_content):
        batch_content_entities = [self.content_dict[c] for c in batch_content]
        batch_content_eid = th.LongTensor(batch_content_entities)
        batch_content_embed = th.mean(self.entity_embedding(batch_content_eid))

        batch_uikg = [self.user_dict[u] for u in batch_user]
        batch_uikg = dgl.batch(batch_uikg)
        batch_keur = self.kg_encoder(batch_uikg)

        score = torch.matmul(torch.matmul(batch_content_embed, self.attention), batch_keur)

        return score


